import torch

from ai import audo_model
from ai.audo_model.s2t.deepspeech2.deepspeech2 import DeepSpeech2Model
from ai.utils import utils_file, utils_model
from ai.audio_dataset.recognition import data_handler
from project.deepspeech2_demo.runner import AishellRunner
from ai.config.config import GxlNode


def main():
    """"""
    config = GxlNode.get_config_from_yaml('./config.yaml')
    # data_handler.build_tokenizer_by_json('./output/aishell.json', config)
    # data_handler.calculate_fbank_cmvn('./output/aishell.json', config)
    data_dev_loader = data_handler.get_iter_by_json('./output/aishell_dev.json', config)
    data_test_loader = data_handler.get_iter_by_json('./output/aishell_test.json', config)
    data_train_loader = data_handler.get_iter_by_json('./output/aishell_train.json', config)
    model = DeepSpeech2Model(
        feat_size=config.dataset.mel_num,
        dict_size=len(data_handler.load_tokenizer()),
        num_conv_layers=2,
        num_rnn_layers=4,
        rnn_size=1024,
        rnn_direction='forward',
        num_fc_layers=2,
        fc_layers_size_list=None,
        use_gru=False
    )
    optimizer = torch.optim.Adam(
        model.parameters(), lr=1e-4
    )
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=5, T_mult=2)
    runner_man = AishellRunner(
        model=model,
        optim=optimizer,
        loss_f=None,
        train_loader=data_train_loader,
        config=config,
        valid_loader=data_dev_loader,
        scheduler=lr_scheduler,
        multi=False,
        local_rank=0,
        is_class=False,
        device=torch.device('cuda:0')
    )
    runner_man.run()
    # runner_man.calculate_valid_loss()
    # data_handler.calculate_fbank_cmvn('./output/aishell.json', config)
    # dev_dataset = data_handler.get_dataset_by_json('./output/aishell_dev.json', config)
    # dev_iter = data_handler.get_iter_by_json('./output/aishell_dev.json', config)
    # for x,x_lens, y ,y_lens in dev_iter:
    #     print(x.shape)
    #     print(x_lens)
    #     print(y.shape)
    #     print(y_lens)


if __name__ == '__main__':
    main()
